package com.shujia.flink.state

import org.apache.flink.api.common.eventtime.WatermarkStrategy
import org.apache.flink.api.common.functions.{ReduceFunction, RuntimeContext}
import org.apache.flink.api.common.serialization.SimpleStringSchema
import org.apache.flink.api.common.state.{ReducingState, ReducingStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.connector.kafka.source.KafkaSource
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

object Demo5ReducingState {
  def main(args: Array[String]): Unit = {
    /**
     * 实时统计每个班级的平均年龄
     */

    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    val source: KafkaSource[String] = KafkaSource
      .builder[String]
      //kafka 集群列表
      .setBootstrapServers("master:9092,node1:9092,node2:9092")
      //消费的topic
      .setTopics("student")
      //消费者组
      .setGroupId("my-group")
      //读取数据的位置，earliest：从最早读取数据，latest：读取最新数据
      .setStartingOffsets(OffsetsInitializer.latest)
      .setValueOnlyDeserializer(new SimpleStringSchema())
      .build

    //使用kafka 数据源
    val studentDS: DataStream[String] = env
      .fromSource(source, WatermarkStrategy.noWatermarks(), "Kafka Source")

    //取出班级和年龄
    val clazzAndAgeDS: DataStream[(String, Int)] = studentDS.map(stu => {
      val split: Array[String] = stu.split(",")
      val clazz: String = split(4)
      val age: Int = split(2).toInt
      (clazz, age)
    })

    //按照班级分组
    val keyByDS: KeyedStream[(String, Int), String] = clazzAndAgeDS.keyBy(_._1)

    val clazzAvgAgeDS: DataStream[(String, Double)] = keyByDS
      .process(new KeyedProcessFunction[String, (String, Int), (String, Double)] {


        //保存总年龄的状态
        var sumAgeState: ReducingState[Double] = _
        //保存总人数的状态
        var numState: ReducingState[Double] = _

        override def open(parameters: Configuration): Unit = {

          val context: RuntimeContext = getRuntimeContext

          //聚合函数
          val reduceFunction: ReduceFunction[Double] = new ReduceFunction[Double] {
            override def reduce(value1: Double, value2: Double): Double = {
              value1 + value2
            }
          }

          //创建聚合状态的描述对象
          val sumAgeStateDesc = new ReducingStateDescriptor[Double]("sumAge", reduceFunction, classOf[Double])
          //创建聚合状态
          sumAgeState = context.getReducingState(sumAgeStateDesc)


          //创建聚合状态的描述对象
          val numStateDesc = new ReducingStateDescriptor[Double]("num", reduceFunction, classOf[Double])
          //创建聚合状态
          numState = context.getReducingState(numStateDesc)

        }

        override def processElement(kv: (String, Int),
                                    ctx: KeyedProcessFunction[String, (String, Int), (String, Double)]#Context,
                                    out: Collector[(String, Double)]): Unit = {
          val (clazz: String, age: Int) = kv


          //向聚合状态中累加一个值，底层会调用reduceFunction函数进行累加
          sumAgeState.add(age)

          //累加计算人数
          numState.add(1)

          //计算平均年龄
          val avgAge: Double = sumAgeState.get() / numState.get()

          //将结果发生到下游
          out.collect((clazz, avgAge))
        }
      })

    clazzAvgAgeDS.print()

    env.execute()

  }

}
